from langchain.agents import ZeroShotAgent, AgentExecutor
from langchain.tools import tool
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

"""
@tool 装饰器自动把函数变成 Agent 可用的“工具”
description 很重要，Agent 会根据描述来“匹配”和“选择工具”。
工具函数尽量“纯函数、健壮清晰”，参数和返回值要简单。
"""


@tool(
    description="用于查询指定城市的天气。参数:city(必填,城市名称,示例:北京,苏州,上海),返回值:实时温度、湿度、天气状况、风向风力。用户问城市天气的时候必须调用这个工具,不可编造数据")
def WeatherTool(city: str) -> str:

    weather_data = {
        "北京": "温度：20℃，湿度：58%，天气：多云，风向：西风，风力：1级",
        "上海": "温度：24℃，湿度：50%，天气：晴，风向：南风，风力：2级",
        "广州": "温度：30℃，湿度：70%，天气：晴，风向：东风，风力：3级",
        "深圳": "温度：28℃，湿度：60%，天气：晴，风向：南风，风力：2级",
    }

    return weather_data.get(city, "没有该城市的天气信息")


@tool(
    description="用于计算加减乘除运算.参数:expression(必填,数学表达式,示例: 2+4,10-2,3*6,10/2,2+3*7),返回值:计算结果.仅处理数字和+、-、*、/")
def CalculatorTool(expression: str) -> str:
    res = eval(expression)
    return f"计算结果:{expression}={res}"


@tool(
    description="用于查询公司知识库中的答案。参数:question(必填,问题,示例:中国有多少个民族,中国有多少个省级行政区),返回值:答案。用户问公司相关问题时必须调用这个工具,不可编造数据")
def RAGTool(question: str) -> str:
    rag_data = {
        "阅兵展示的武器": "阅兵展示的武器有东风-17、东风-41、东风-50、东风-21D、东风-31A、东风-41A、东风-41B、东风-41C、东风-41D、东风-41E、东风-41F、东风-41G、东风-41H、东风-41J、东风-41K、东风-41L、东风-41M、东风-41N、东风-41O、东风-41P、东风-41Q、东风-41R、东风-41S、东风-41T、东风-41U、东风-41V、东风-41W、东风-41X、东风-41Y、东风-41Z、东风-41AA、东风-41AB、东风-41AC、东风-41AD、东风-41AE",
        "中国的领土面积": "中国的领土面积960万平方公里",
        "中国有多少个民族": "中国有56个民族",
        "中国有多少个省级行政区": "中国有34个省级行政区"
    }

    for key, value in rag_data.items():
        if key in question:
            return value
    return f"在知识库中没有找到[{question}]答案"


tools = [WeatherTool, CalculatorTool, RAGTool]

llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key="sk-6434b554122c4b049ceb805d703f695b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)
prompt_template = """
你是一个帮助用户解决问题的智能助手，你拥有以下工具:
{tools}

使用规则:
1:先理解问题:无需工具则直接回答,需要工具则调用工具(根据工具的描述匹配)
2:调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全）；
3:拿到工具结果后, 判断是否继续调用工具, 如果需要继续调用工具,则调用工具,否则整理,返回结果;


用户的问题:{input}
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "input"],
)

#构建Agent
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True,
)

#构建执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
)

if __name__ == '__main__':
    res1 = agent_executor.invoke({"input":"什么是人工智能, 请简单的解释一下"})
    print(res1["output"])

    # res2 = agent_executor.invoke({"input": "查询北京今天的天气"})
    # print(res2["output"])

    #res3 = agent_executor.invoke({"input":"根据上海的天气, 推荐合适的穿搭,并计算这套穿搭需要多少钱(裤衩99元,背心:87,帽子100)"})
    #print(res3["output"])

    #res4 = agent_executor.invoke({"input":"中国有多少个民族"})
    #print(res4["output"])

